#!/usr/bin/python3
# -*- coding:utf-8 -*-
# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================

import os
import torch
import numpy as np

def rwkv_time_mix(
    B, H, T, N,
    k: np.ndarray,
    v: np.ndarray,
    w: np.ndarray,
    q: np.ndarray,
    a: np.ndarray,
    b: np.ndarray,
    h: np.ndarray):
    # N = C // H  # 头的维度
    param_shape = (B, T, H, N)
    u_shape = (H, N)
    state_shape = (B, H, N, N)
   
    o = np.zeros(param_shape).astype(np.float16)

    for t in range(T):
        for bi in range(B):
            for hi in range(H):
                q_t = q[bi, t, hi]
                k_t = k[bi, t, hi]
                v_t = v[bi, t, hi]
                a_t = a[bi, t, hi]
                b_t = b[bi, t, hi]
                w_t = np.exp(w[bi, t, hi])

                sa = np.sum((a_t[None, :] * h[bi, hi]), axis=1)
                h[bi, hi] = (h[bi, hi] * w_t[None, :] +  # [N,V] * [N,1]
                                 k_t[None, :] * v_t[:, None] +     # [N,1] * [1,V]
                                 sa[:, None] * b_t[None, :])                # [V] * [N,1]

                y = np.sum((h[bi, hi] * q_t[None, :]), axis=1)

                o[bi, t, hi] = y
    return o, h

if __name__ == "__main__":
    B, T, H, N = 1, 256, 40, 64
    cur_dir = os.path.dirname(os.path.abspath(__file__))
    np.set_printoptions(threshold=np.inf)
    data_type = torch.float16

    k = torch.randn(B, T, H, N).to(data_type)
    v = torch.randn(B, T, H, N).to(data_type)
    w = torch.randn(B, T, H, N).to(data_type)
    r = torch.randn(B, T, H, N).to(data_type)
    a = torch.randn(B, T, H, N).to(data_type)
    b = torch.randn(B, T, H, N).to(data_type)
    h0 = torch.randn(B, H, N, N).to(data_type)

    input_dir = "./input/"
    k.numpy().tofile(os.path.join(input_dir,  "input_k.bin"))
    v.numpy().tofile(os.path.join(input_dir,  "input_v.bin"))
    w.numpy().tofile(os.path.join(input_dir,  "input_w.bin"))
    r.numpy().tofile(os.path.join(input_dir,  "input_r.bin"))
    a.numpy().tofile(os.path.join(input_dir,  "input_a.bin"))
    b.numpy().tofile(os.path.join(input_dir,  "input_b.bin"))
    h0.numpy().tofile(os.path.join(input_dir,  "input_h0.bin"))

    out, ht = rwkv_time_mix(B, H, T, N, k.numpy(), v.numpy(), w.numpy(),
                                r.numpy(), a.numpy(), b.numpy(),
                                h0.numpy())
    print(out.astype(np.float16)[0][0][1][:10])
    
    out.astype(np.float16).tofile(os.path.join( "./output/output_o_golden.bin"))
    ht.astype(np.float16).tofile(os.path.join( "./output/output_ht_golden.bin"))